import json, asyncio
import uuid
from fastapi import FastAPI, APIRouter, WebSocket, WebSocketDisconnect
from typing import List, Dict, Set
from src.common.response import success_response
from src.core.session import SQLCON
from src.entity.models import UpdateMessage, Message
from src.servers.message import msg_query, msg_create, msg_update
from src.core.logconfig import log
from src.utils.tools import get_cur_time

ws = APIRouter()


class ConnectionManager:

    def __init__(self):
        self.active_conns: Dict[str, Dict[str, WebSocket]] = {}
        self.lock = asyncio.Lock()

    # 连接websocket
    async def connect(self, websocket: WebSocket, room_id: str, user_id: str):
        # log.logger.info(f"websocket connected! {room_id} -- {user_id}")
        await websocket.accept()
        async with self.lock:
            if room_id not in self.active_conns:
                self.active_conns[room_id] = {}
            self.active_conns[room_id][user_id] = websocket
        return user_id

    # 断开websocket连接
    async def disconnect(self, room_id:str, user_id:str):
        async with self.lock:
            if room_id in self.active_conns and user_id in self.active_conns[room_id]:
                del self.active_conns[room_id][user_id]
                if not self.active_conns[room_id]:
                    del self.active_conns[room_id] 
        
    # 发送消息
    # async def send_message(self, msg: dict, websocket: WebSocket):
    #     log.logger.info(f"send msg:{ msg }")
    #     await websocket.send_json(msg)
    async def broadcast(self, room_id:str, message: dict, exclude_user: str = None):
        if room_id not in self.active_conns:
            return
        # log.logger.info(f"websocket broadcast! room: {room_id} message: {message} ")
        async with self.lock:
            for user_id, websocket in self.active_conns[room_id].items():
                if user_id != exclude_user:
                    try:
                        await websocket.send_text(json.dumps(message))
                    except Exception as e:
                        print(f"Error sending message to {user_id}: {e}")
                        await self.disconnect(room_id, user_id)
                        

    # 将消息发送到活跃的连接
    # async def broadcast2(self, msg: dict):
    #     for connect in self.active_conns:
    #         log.logger.info(f"broadcast msg { msg }")
    #         await connect.send_json(msg)


manager = ConnectionManager()


# 生成唯一用户ID
def generate_user_id(u:str=None):
    if u is None:
        return str(uuid.uuid4())[:8]
    else:
        return str(u)
    
def receive_data(user, msg, msg_type:int=0, editable:bool=False):
    return {
        "user_id": user,
        "content": msg,
        "username": user,
        "msg_type": msg_type,
        "editable": editable,
        "create_time": get_cur_time(),
    }

@ws.post("/ws/update")
async def edit_msg(msg: UpdateMessage):
    resp = SQLCON.db_query_by_id(msg.id)
    if resp:
        success_response(SQLCON.db_update(entity=Message, update=msg))
    else:
        success_response("未查询到数据")

@ws.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket, u=None):
    user_id = generate_user_id(u)
    await manager.connect(websocket, room_id='chat', user_id=user_id)
    await manager.broadcast("chat", receive_data(user_id, f"Welcome {user_id} into chat", msg_type=0))

    try:
        while True:
            data = await websocket.receive_text()
            msg = receive_data(user_id, json.loads(data)['content'])
            await msg_create(msg)
            await manager.broadcast('chat', msg )
    except WebSocketDisconnect:
        manager.disconnect(websocket, user_id)
        await manager.broadcast(receive_data(user_id, "leaved"))
        

# WebSocket 实时协作端点
@ws.websocket("/ws/{room_id}/{user_id}")
async def websocket_endpoint2(websocket: WebSocket, room_id:str, user_id:str, msg_type):
    await manager.connect(websocket, room_id, user_id)
    try:
        while True:
            data = await websocket.receive_text()
            # 这里可以添加广播逻辑，将更新发送给所有连接的客户端
            message = receive_data(user_id, json.loads(data)['content'], msg_type=msg_type)
            if message['msg_type'] == '0':
                # 广播聊天消息
                chat_msg = Message(
                    content=message['text'],
                    user_id=user_id,
                    username=message.get('username', 'Anonymous')
                )
                await manager.broadcast(room_id, {
                        'type': '0',
                        'message': chat_msg.dict()
                    }
                )
            elif message['msg_type'] == '1':
                
                # 广播思维导图更新
                await manager.broadcast(room_id, {
                    "type": '1',
                    'content': message['content'],
                    'user_id': user_id
                }, exclude_user=user_id)
    except Exception as e:
        print(f"Websocket error: {e}")
    finally:
        await manager.disconnect(room_id, user_id)